from flask import Flask, request, jsonify, render_template, redirect, url_for, flash
import torch
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import scipy.io as sio
import time
import spectral
from utils import HybridSN, applyPCA, padWithZeros
import os
from concurrent.futures import ThreadPoolExecutor

# 设置Matplotlib使用非交互式后端
matplotlib.use('Agg')

# 初始化Flask应用
app = Flask(__name__)
app.secret_key = 'your_secret_key_here'  # 设置密钥

# 选择设备（GPU或CPU）
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 加载模型
model = HybridSN().to(device)
model.load_state_dict(torch.load(os.path.join('model.pth'), weights_only=True))
model.eval()
# 配置文件上传的保存路径
UPLOAD_FOLDER = 'uploads'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

# 创建线程池（用于并行逐帧预测）
executor = ThreadPoolExecutor(max_workers=8)

# 主路由处理分类请求
@app.route('/', methods=["GET", "POST"])
def classify():
    if request.method == 'GET':
        return render_template('index.html')
    elif request.method == 'POST':
        # 获取上传的文件
        file1 = request.files['image1']
        file2 = request.files['image2']

        if file1.filename == '' or file2.filename == '':
            flash("未选择文件或文件名为空")
            return redirect(url_for('classify'))

        # 保存文件到服务器
        file1_path = os.path.join(app.config['UPLOAD_FOLDER'], file1.filename)
        file2_path = os.path.join(app.config['UPLOAD_FOLDER'], file2.filename)
        file1.save(file1_path)
        file2.save(file2_path)

        flash("文件已成功上传并保存到本地！")
        print("文件已成功上传并保存到本地！")

        # 处理分类数据
        return process_classification(file1_path, file2_path)

# 逐帧预测的函数
def process_patch(i, X, y, patch_size):
    width = y.shape[1]
    row_output = np.zeros((width,))
    for j in range(width):
        if int(y[i, j]) == 0:
            continue
        image_patch = X[i:i + patch_size, j:j + patch_size, :]
        image_patch = image_patch.reshape(1, image_patch.shape[0], image_patch.shape[1], image_patch.shape[2], 1)
        X_test_image = torch.FloatTensor(image_patch.transpose(0, 4, 3, 1, 2)).to(device)
        with torch.no_grad():
            prediction = model(X_test_image)
        row_output[j] = np.argmax(prediction.detach().cpu().numpy()) + 1
    return i, row_output

# 并行处理逐行预测
def parallel_prediction(X, y, patch_size):
    height = y.shape[0]
    outputs = np.zeros_like(y)
    futures = [executor.submit(process_patch, i, X, y, patch_size) for i in range(height)]
    for future in futures:
        i, row_output = future.result()
        outputs[i, :] = row_output
        if i % 20 == 0:
            print(f'... 正在处理第 {i} 行 ...')
    return outputs

# 分类数据处理函数
def process_classification(file1_path, file2_path):
    try:
        # 加载文件1（X数据）
        mat_data = sio.loadmat(file1_path)
        X = None
        for key in mat_data.keys():
            value = mat_data[key]
            if key not in ['__header__', '__version__', '__globals__']:
                if isinstance(value, np.ndarray) and len(value.shape) == 3:
                    X = value
                    break
        if X is None:
            raise ValueError("没有找到有效的X数据")

        # 加载文件2（y数据）
        mat_data1 = sio.loadmat(file2_path)
        y = None
        for key in mat_data1.keys():
            value = mat_data1[key]
            if key not in ['__header__', '__version__', '__globals__']:
                if isinstance(value, np.ndarray) and len(value.shape) == 2:
                    y = value
                    break
        if y is None:
            raise ValueError("没有找到有效的y数据")

        # 图像维度信息
        height, width = y.shape
        patch_size = 25
        pca_components = 30

        # PCA降维和零填充
        X = applyPCA(X, numComponents=pca_components)
        X = padWithZeros(X, patch_size // 2)

        # 并行逐行预测
        outputs = parallel_prediction(X, y, patch_size)

        # 输出信息
        print(f"输出图像形状: {outputs.shape}")
        print(f"输出图像最小值: {outputs.min()}, 最大值: {outputs.max()}")

        # 使用spectral库显示图像
        spectral.imshow(classes=outputs.astype(int), figsize=(5, 5))

        # 生成图像文件
        output_image_path = os.path.join('static', f'output_image_{int(time.time())}.png')
        plt.savefig(output_image_path)
        plt.close()

        # 返回图像URL
        return jsonify({'image_url': f'/static/{os.path.basename(output_image_path)}'})

    except Exception as e:
        print(f"分类过程中发生错误: {e}")
        return jsonify({'error': str(e)})

# 提供静态文件
@app.route('/static/<path:path>')
def serve_static(path):
    return app.send_static_file(path)

# 启动Flask应用
if __name__ == '__main__':
    app.run(threaded=True, host='0.0.0.0', port=5000)
